package install.ren.exercise.test.zookeeper;

import lombok.Data;
import lombok.SneakyThrows;
import lombok.extern.slf4j.Slf4j;
import org.apache.curator.framework.CuratorFramework;
import org.apache.curator.framework.CuratorFrameworkFactory;
import org.apache.curator.retry.ExponentialBackoffRetry;

import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * @author jonny
 * 计算nodePath下所有非叶子节点的子节点个数
 * Created on 2021/12/15
 */
@Slf4j
public class ZookeeperNodeCounter {
    public static String connectString = "?:2181";
    public static List<NodeItem> NODES = new ArrayList<>();
    public static AtomicInteger SUM = new AtomicInteger(0);

    public static void main(String[] args) {
        try {
            String nodePath = "/";
            CuratorFramework client = CuratorFrameworkFactory.newClient(connectString, new ExponentialBackoffRetry(20000, 3));
            client.start();
            NodeItem parent = new NodeItem(nodePath, new ArrayList<>());
            NODES.add(parent);
            final List<String> strings = client.getChildren().forPath("/");
            strings.forEach(p -> {
                getNode(parent, "/".concat(p), client);
            });


            System.out.println("cnt:" + SUM.get());
            Collections.sort(NODES);
            NODES.forEach((k) -> {
                if (k.getChildCnt() > 10) {
                    System.out.println(k.getChildCnt() + ":" + k.getPath());
                }
            });

        } catch (Exception e) {
            e.printStackTrace();
        }
    }

    @SneakyThrows
    public static void getNode(NodeItem parentNode, String path, CuratorFramework client) {
        SUM.incrementAndGet();

        final List<String> strings = client.getChildren().forPath(path);
        if (strings != null && strings.size() > 0) {

            strings.forEach(p -> {
                try {
                    String subPath = path + "/" + p;
                    NodeItem nodeItem = new NodeItem(subPath, new ArrayList<>());
                    NODES.add(nodeItem);
                    parentNode.add(nodeItem);
                    getNode(nodeItem, subPath, client);
                } catch (Exception e) {
                    e.printStackTrace();
                }
            });

        } else {
            NodeItem nodeItem = new NodeItem( path, new ArrayList<>());
            parentNode.add(nodeItem);
        }
    }

    @Data
    static class NodeItem implements Comparable<NodeItem> {
        private String path;
        private List<NodeItem> nodeItems;
        private Integer total = 1;

        public NodeItem(String path, List<NodeItem> nodeItems) {
            this.path = path;
            this.nodeItems = nodeItems;
        }

        public void add(NodeItem nodeItem) {
            nodeItems.add(nodeItem);
        }

        @Override
        public int compareTo(NodeItem o) {
            return this.getChildCnt().compareTo(o.getChildCnt());
        }

        private Integer getChildCnt() {
            if (nodeItems.size() != 0) {
                if (total == 1) {
                    total = nodeItems.stream().mapToInt(NodeItem::getChildCnt).sum();
                    total++;
                }
                return total;
            }else {
                return 0;
            }
        }
    }


}
